import glob
import torchvision.transforms as transforms
import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image
import random


class CustomDataset(Dataset):
    def __init__(self, image_dir, label_dir, img_format, lab_format, task='train', size = (256, 512)):
        images = glob.glob(r'./data/{}/{}/*/*{}'.format(image_dir, task, img_format))
        labels = glob.glob(r'./data/{}/{}/*/*{}'.format(label_dir, task, lab_format))
        images.sort()
        labels.sort()
        self.data = list(zip(images, labels))
        self.transforms1 = transforms.Compose([transforms.Resize(size, Image.BILINEAR),
                                              transforms.ToTensor(),
                                               transforms.Normalize((.485, .456, .406), (.229, .224, .225)),])
        self.transforms2 = transforms.Resize(size, Image.NEAREST)

    def __getitem__(self, index):
        image_path, label_path = self.data[index]
        image = Image.open(image_path)
        label = Image.open(label_path)
        image, label = self.__randomtransform__(image, label)
        image = self.transforms1(image)
        label = torch.LongTensor(np.array(self.transforms2(label)).astype('int32'))
        return image, label, image_path.rsplit('/',1)[1]

    def __len__(self):
        return len(self.data)

    def __randomtransform__(self, img, label):
        if random.random()<0.5:
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
            label = label.transpose(Image.FLIP_LEFT_RIGHT)
        if random.random()<0.5:
            img = img.transpose(Image.FLIP_TOP_BOTTOM)
            label = label.transpose(Image.FLIP_LEFT_RIGHT)
        return img, label

class TestDataset(Dataset):
    def __init__(self, image_dir, img_format, size = (256, 512)):
        self.images = glob.glob(r'./data/{}/*{}'.format(image_dir, img_format))
        self.transforms1 = transforms.Compose([transforms.Resize(size, Image.BILINEAR),
                                              transforms.ToTensor(),
                                               transforms.Normalize((.485, .456, .406), (.229, .224, .225)),])

    def __getitem__(self, index):
        image_path = self.images[index]
        image = Image.open(image_path)
        image = self.__randomtransform__(image)
        image = self.transforms1(image)
        return image, image_path.rsplit('/',1)[1]

    def __len__(self):
        return len(self.images)

    def __randomtransform__(self, img):
        if random.random()<0.5:
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
        if random.random()<0.5:
            img = img.transpose(Image.FLIP_TOP_BOTTOM)
        return img